import os
import json
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import xml.etree.ElementTree as ET
from glob import glob
from scipy.ndimage import gaussian_filter

def json_points(label):
    with open(label, 'r') as f:
        info = json.load(f)
    points = info['points']
    return points

def xml_points(label):
    tree = ET.parse(label)
    root = tree.getroot()
    points = [[int(obj.find('point').find('x').text), int(obj.find('point').find('y').text)]for obj in root.findall('object')]
    return points
    
def get_result(pt_file, shape):
    density = torch.load(pt_file)
    count = "{:.2f}".format(density.sum().item())

    resize_density = F.interpolate(density, size=shape, mode="bilinear")
    
    x_sum = torch.sum(density, dim=(-1, -2))
    scale_factor = torch.nan_to_num(torch.sum(resize_density, dim=(-1, -2)) / x_sum, nan=0.0, posinf=0.0, neginf=0.0)
    resize_density = resize_density * scale_factor
    
    return count, resize_density.cpu().squeeze()

def generate_density_map(shape, label, sigma = 8):
    """
    Generate the density map based on the dot annotations provided by the label.
    """
    height, width = shape
    label = torch.tensor(label).float()
    
    density_map = torch.zeros((1, height, width), dtype=torch.float32)

    if len(label) > 0:
        assert len(label.shape) == 2 and label.shape[1] == 2, f"label should be a Nx2 tensor, got {label.shape}."
        label_ = label.long()
        label_[:, 0] = label_[:, 0].clamp(min=0, max=width - 1)
        label_[:, 1] = label_[:, 1].clamp(min=0, max=height - 1)
        density_map[0, label_[:, 1], label_[:, 0]] = 1.0

    if sigma is not None:
        assert sigma > 0, f"sigma should be positive if not None, got {sigma}."
        density_map = torch.from_numpy(gaussian_filter(density_map, sigma=sigma))

    return density_map.squeeze()

def LMDS_counting(input):
    input_max = torch.max(input).item()

    keep = nn.functional.max_pool2d(input, (3, 3), stride=1, padding=1)
    keep = (keep == input).float()
    input = keep * input

    '''set the pixel valur of local maxima as 1 for counting'''
    input[input < 100.0 / 255.0 * input_max] = 0
    input[input > 0] = 1

    ''' negative sample'''
    if input_max < 0.1:
        input = input * 0

    count = int(torch.sum(input).item())

    return count




savepath = "./02_comp_rgb_rgbtcc"
if not os.path.exists(savepath):
    os.makedirs(savepath)

gt_root = '/data/store1/nzd/tir_cc/datasets/rgbtcc/test'
image_root = '/data/store1/nzd/tir_cc/datasets/rgbtcc/test'

results = {
    'CLIP-EBC': '/data/store1/nzd/tir_cc/methods/clip/1_adapt/rgbtcc_results',
    'BL': '/data/store1/nzd/tir_cc/methods/bayesian/1_adapt_rgbtcc/results',
    'P2P': '/data/store1/nzd/tir_cc/methods/p2p/p2p_rgbtcc/results',
    'HMoDE': '/data/store1/nzd/tir_cc/methods/hmode/RMS_rgbtcc/results',
    'STEERER': '/data/store1/nzd/tir_cc/methods/steerer/STEERER_rgbtcc/results',
    'FIDTM': '/data/store1/nzd/tir_cc/methods/fidtm/25_optim_rgbtcc/results',
}

gt_files = glob(os.path.join(gt_root, '*_GT.json'))
for gt_file in gt_files:
    name = gt_file.split('/')[-1].split('.')[0].replace('_GT', 'R')

    img_file = gt_file.replace('_GT.json', '_T.jpg')
    img = plt.imread(img_file)
    shape = img.shape[:2]

    if shape == (640, 480):
        continue
        
    points = json_points(gt_file)
    gt_num = len(points)
    gt_density = generate_density_map(shape, points, sigma=8)

    # 保存对比结果
    res_num = len(results)
    fig, axes = plt.subplots(1, res_num + 2, dpi=800, tight_layout=True, frameon=False)
    
    axes[0].imshow(img)
    if len(points) > 0:
        axes[0].scatter([p[0] for p in points], [p[1] for p in points], s=0.8, c="red", marker=".", edgecolors="none")
    axes[0].axis("off")
    axes[0].set_title(f"{name}", fontsize=6)

    axes[1].imshow(img, cmap="jet")
    axes[1].imshow(gt_density, cmap="jet", alpha=0.7)
    axes[1].axis("off")
    axes[1].set_title(f"GT: {gt_num}", fontsize=6)

    i = 1
    for method, path in results.items():
        if method == 'P2P':
            pre_points = torch.load(os.path.join(path, f'{name}.pt')).detach().cpu().numpy().tolist()
            pre_count = "{:.2f}".format(len(pre_points))
            pre_density = generate_density_map(shape, pre_points, sigma=8)
        elif method == 'FIDTM':
            _, pre_density = get_result(os.path.join(path, f'{name}.pt'), shape)
            pre_count = LMDS_counting(torch.load(os.path.join(path, f'{name}.pt')))
            pre_count = "{:.2f}".format(pre_count)
        else:
            pre_count, pre_density = get_result(os.path.join(path, f'{name}.pt'), shape)

        if method == 'HMoDE':
            pre_count = "{:.2f}".format(float(pre_count) / 100.)
        
        axes[i+1].imshow(img, cmap="jet")
        axes[i+1].imshow(pre_density, cmap="jet", alpha=0.7)
        axes[i+1].axis("off")
        axes[i+1].set_title(f"{method}: {pre_count}", fontsize=6)

        i += 1
    
    plt.savefig(f"{savepath}/{name}", bbox_inches='tight', pad_inches=0)
    # plt.show()
    plt.close()




